from shared_functions import *

def is_weekend(tx_datetime):
    """
    判断交易是否发生在周末
    
    参数:
        tx_datetime: 交易日期时间
        
    返回:
        int: 0表示工作日，1表示周末
    """
    
    # 将日期转换为星期几（0是周一，6是周日）
    weekday = tx_datetime.weekday()
    # 二进制值：工作日为0，周末为1
    is_weekend = weekday>=5
    
    return int(is_weekend)

def is_night(tx_datetime):
    """
    判断交易是否发生在夜间
    
    参数:
        tx_datetime: 交易日期时间
        
    返回:
        int: 1表示凌晨6点前（夜间），0表示其他时间
    """
    
    # 获取交易的小时数
    tx_hour = tx_datetime.hour
    # 二进制值：如果小时数小于等于6则为1，否则为0
    is_night = tx_hour<=6
    
    return int(is_night)

def get_customer_spending_behaviour_features(customer_transactions, windows_size_in_days=[1,7,30]):
    """
    获取客户消费行为特征
    
    参数:
        customer_transactions: 客户交易数据
        windows_size_in_days: 时间窗口大小（天）
        
    返回:
        包含新特征的客户交易数据
    """
    
    # 首先按时间顺序排列交易记录
    customer_transactions=customer_transactions.sort_values('TX_DATETIME')
    
    # 将交易日期时间设置为索引，这将允许使用滚动函数
    customer_transactions.index=customer_transactions.TX_DATETIME
    
    # 对于每个时间窗口大小
    for window_size in windows_size_in_days:
        
        # 计算给定时间窗口内交易金额的总和和交易次数
        SUM_AMOUNT_TX_WINDOW=customer_transactions['TX_AMOUNT'].rolling(str(window_size)+'d').sum()
        NB_TX_WINDOW=customer_transactions['TX_AMOUNT'].rolling(str(window_size)+'d').count()
    
        # 计算给定时间窗口内的平均交易金额
        # NB_TX_WINDOW总是大于0，因为当前交易总是包含在内
        AVG_AMOUNT_TX_WINDOW=SUM_AMOUNT_TX_WINDOW/NB_TX_WINDOW
    
        # 保存特征值
        customer_transactions['CUSTOMER_ID_NB_TX_'+str(window_size)+'DAY_WINDOW']=list(NB_TX_WINDOW)
        customer_transactions['CUSTOMER_ID_AVG_AMOUNT_'+str(window_size)+'DAY_WINDOW']=list(AVG_AMOUNT_TX_WINDOW)
    
    # 根据交易ID重新索引
    customer_transactions.index=customer_transactions.TRANSACTION_ID
        
    # 返回包含新特征的DataFrame
    return customer_transactions



def get_count_risk_rolling_window(terminal_transactions, delay_period=7, windows_size_in_days=[1,7,30], feature="TERMINAL_ID"):
    """
    获取终端风险滚动窗口统计特征
    
    参数:
        terminal_transactions: 终端交易数据
        delay_period: 延迟期（天）
        windows_size_in_days: 时间窗口大小（天）
        feature: 特征名称
        
    返回:
        包含风险特征的终端交易数据
    """
    
    # 按时间顺序排列交易记录
    terminal_transactions=terminal_transactions.sort_values('TX_DATETIME')
    
    # 将交易日期时间设置为索引
    terminal_transactions.index=terminal_transactions.TX_DATETIME
    
    # 计算延迟期内的欺诈交易数和总交易数
    NB_FRAUD_DELAY=terminal_transactions['TX_FRAUD'].rolling(str(delay_period)+'d').sum()
    NB_TX_DELAY=terminal_transactions['TX_FRAUD'].rolling(str(delay_period)+'d').count()
    
    # 对于每个时间窗口大小
    for window_size in windows_size_in_days:
    
        # 计算(延迟期+窗口期)内的欺诈交易数和总交易数
        NB_FRAUD_DELAY_WINDOW=terminal_transactions['TX_FRAUD'].rolling(str(delay_period+window_size)+'d').sum()
        NB_TX_DELAY_WINDOW=terminal_transactions['TX_FRAUD'].rolling(str(delay_period+window_size)+'d').count()
    
        # 计算窗口期内的欺诈交易数和总交易数
        NB_FRAUD_WINDOW=NB_FRAUD_DELAY_WINDOW-NB_FRAUD_DELAY
        NB_TX_WINDOW=NB_TX_DELAY_WINDOW-NB_TX_DELAY
    
        # 计算窗口期内的风险评分
        RISK_WINDOW=NB_FRAUD_WINDOW/NB_TX_WINDOW
        
        # 保存特征值
        terminal_transactions[feature+'_NB_TX_'+str(window_size)+'DAY_WINDOW']=list(NB_TX_WINDOW)
        terminal_transactions[feature+'_RISK_'+str(window_size)+'DAY_WINDOW']=list(RISK_WINDOW)
        
    # 根据交易ID重新索引
    terminal_transactions.index=terminal_transactions.TRANSACTION_ID
    
    # 用0替换NA值（所有NB_TX_WINDOW为0时的未定义风险评分）
    terminal_transactions.fillna(0,inplace=True)
    
    return terminal_transactions


# 输入目录
DIR_INPUT='./simulated-data-raw/' 

# 开始日期
BEGIN_DATE = "2025-04-01"
# 结束日期
END_DATE = "2025-09-30"

print("加载文件")
# 从文件中读取交易数据
transactions_df=read_from_files(DIR_INPUT, BEGIN_DATE, END_DATE)
print("已加载{0}笔交易，其中包含{1}笔欺诈交易".format(len(transactions_df),transactions_df.TX_FRAUD.sum()))

# 添加交易时间特征
transactions_df['TX_DURING_WEEKEND']=transactions_df.TX_DATETIME.apply(is_weekend)
transactions_df['TX_DURING_NIGHT']=transactions_df.TX_DATETIME.apply(is_night)


# 获取客户ID为0的消费行为特征示例
spending_behaviour_customer_0=get_customer_spending_behaviour_features(transactions_df[transactions_df.CUSTOMER_ID==0])

# 对所有客户应用消费行为特征计算
transactions_df=transactions_df.groupby('CUSTOMER_ID').apply(lambda x: get_customer_spending_behaviour_features(x, windows_size_in_days=[1,7,30]))
transactions_df=transactions_df.sort_values('TX_DATETIME').reset_index(drop=True)

# 对所有终端应用风险滚动窗口特征计算
transactions_df=transactions_df.groupby('TERMINAL_ID').apply(lambda x: get_count_risk_rolling_window(x, delay_period=7, windows_size_in_days=[1,7,30], feature="TERMINAL_ID"))
transactions_df=transactions_df.sort_values('TX_DATETIME').reset_index(drop=True)


# 输出目录
DIR_OUTPUT = "./simulated-data-transformed/"

# 如果输出目录不存在则创建
if not os.path.exists(DIR_OUTPUT):
    os.makedirs(DIR_OUTPUT)

# 设置起始日期
start_date = datetime.datetime.strptime("2025-04-01", "%Y-%m-%d")

# 按天分割数据并保存为pickle文件
for day in range(transactions_df.TX_TIME_DAYS.max()+1):
    
    # 获取当天的交易数据
    transactions_day = transactions_df[transactions_df.TX_TIME_DAYS==day].sort_values('TX_TIME_SECONDS')
    
    # 计算日期
    date = start_date + datetime.timedelta(days=day)
    filename_output = date.strftime("%Y-%m-%d")+'.pkl'
    
    # Protocol=4 适用于Google Colab
    transactions_day.to_pickle(DIR_OUTPUT+filename_output, protocol=4)